﻿using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;
using System.Text;

namespace Elderos.TextSearchers
{
	public class AhoCorasickTokenizedSearcher : TokenizedTextSearcherBase
	{
		public AhoCorasickTokenizedSearcher(IStringLoader stringLoader, ITokenizer tokenizer) : 
			base(stringLoader, tokenizer) 
		{}

		private Automaton _automaton = new Automaton();

		#region Overrides of TextSearcherBase

		private int _loadedStringsCount;
		public override int LoadedStringsCount
		{
			get { return _loadedStringsCount; }
		}

		protected override IEnumerable<Position> InnerSearch(string[] tokens)
		{
			var results = new List<Position>();
			Node currentNode = _automaton.Root;
			for (int i = 0; i < tokens.Length; i++)
			{
				Node nextNode;
				do
				{
					nextNode = _automaton.Goto(currentNode, tokens[i]);
					if (nextNode == null)
						currentNode = _automaton.Fail(currentNode);
				} while (nextNode == null);

				currentNode = nextNode;

				IEnumerable<int> nodeResults = _automaton.GetResults(currentNode);

				if (nodeResults != null)
					results.AddRange(nodeResults.Select(length => new Position(i - length + 1, length)));
			}

			return results;
		}

		protected override void ProcessString(string str)
		{
			string[] tokens = Tokenizer.Tokenize(str);

			//ignoring untokenizable strings
			if (tokens == null || tokens.Length <= 0)
				return;

			_loadedStringsCount++;

			Node currentNode = _automaton.Root;
			for (int i = 0; i < tokens.Length; i++)
			{
				//interning strings to lower memory consumption
				string currentToken = String.Intern(tokens[i]);

				if (currentNode.Children.ContainsKey(currentToken))
				{
					currentNode = currentNode.Children[currentToken];
					if (i == tokens.Length - 1)
					{
						if (_automaton.ResultsDict.ContainsKey(currentNode))
							_automaton.ResultsDict[currentNode].Add(tokens.Length);
						else
							_automaton.ResultsDict.Add(currentNode, new List<int> { tokens.Length });
					}
					continue;
				}

				var nextNode = new Node();
				currentNode.Children.Add(currentToken, nextNode);
				if (i == tokens.Length - 1)
				{
					if (_automaton.ResultsDict.ContainsKey(nextNode))
						_automaton.ResultsDict[nextNode].Add(tokens.Length);
					else
						_automaton.ResultsDict.Add(nextNode, new List<int> { tokens.Length });
				}

				currentNode = nextNode;
			}
		}

		protected override void OnStringsLoaded()
		{
			base.OnStringsLoaded();
			_automaton.ProcessNodeFailures();
		}

		#endregion

		private class Automaton
		{
			public Automaton()
			{
				_failureDict.Add(_root, _root);
			}

			private Node _root = new Node();
			public Node Root
			{
				get { return _root; }
			}

			private Dictionary<Node, Node> _failureDict = new Dictionary<Node, Node>();
			public Dictionary<Node, Node> FailureDict
			{
				get { return _failureDict; }
			}

			private Dictionary<Node, List<int>> _resultsDict = new Dictionary<Node, List<int>>();
			public Dictionary<Node, List<int>> ResultsDict
			{
				get { return _resultsDict; }
			}

			public void ProcessNodeFailures()
			{
				var queue = new Queue<Node>();
				foreach (var node in _root.Children.Values)
				{
					queue.Enqueue(node);
				}

				while (queue.Count > 0)
				{
					Node node = queue.Dequeue();
					foreach (var kv in node.Children)
					{
						Node child = kv.Value;
						string token = kv.Key;

						queue.Enqueue(child);
						Node failure = Fail(child);
						while (Goto(failure, token) == null)
							failure = Fail(failure);

						Node failureDestination = Goto(failure, token);

						if (failureDestination != _root)
							_failureDict.Add(child, failureDestination);

						if (_resultsDict.ContainsKey(failure))
						{
							if (!_resultsDict.ContainsKey(child))
								_resultsDict.Add(child, new List<int>());
							_resultsDict[child].AddRange(_resultsDict[failure]);
						}

					}
				}
			}

			public Node Goto(Node fromNode, string token)
			{
				Node outNode;
				if (fromNode.Children.TryGetValue(token, out outNode))
					return outNode;
				if (fromNode == _root)
					return _root;
				return null;
			}

			public Node Fail(Node fromNode)
			{
				Node outNode;
				if (FailureDict.TryGetValue(fromNode, out outNode))
					return outNode;
				return _root;
			}

			public IEnumerable<int> GetResults(Node node)
			{
				List<int> lengths;
				if (ResultsDict.TryGetValue(node, out lengths))
					return lengths;
				return null;
			}
		}

		private class Node
		{
			public Node()
			{
				Children = new Dictionary<string, Node>();
			}

			public Dictionary<string, Node> Children { get; private set; }

			public override string ToString()
			{
				return "Node: " + GetHashCode() % 1000;
			}
		}
	}
}
